from __future__ import division
from xml.etree import cElementTree
import sys, re, random
from numpy import empty, dot, sum
from util import *

ALPHA = 0.1
ITERATIONS = 128

class PageRank:
    def __init__(self, file):
        self._file = file
        self.titles = dict()
        self.transitions = None
        
        self.gettitles()
        self.gettransitions()
        self.pagerank()

    @property
    def docnum(self):
        return len(self.titles)
    
    def gettitles(self):
        '''Store titles and docids.'''
        title = ''
        docid = 0
        for event, elem in cElementTree.iterparse(self._file):
            t = re.sub('{.*}','', elem.tag.lower())
            if t == 'page':
                self.titles[title.lower()] = docid
                docid += 1
            elif t == 'title':
                title = elem.text
            elem.clear()

    def gettransitions(self):
        N = self.docnum
        self.transitions = empty((N, N))
        docid = 0
        alphabet = set()
        ones = 0
        for event, elem in cElementTree.iterparse(self._file):
            t = re.sub('{.*}','', elem.tag.lower())
            if t == 'page':
                docid += 1
            elif t == 'text':
                links = filter(lambda x: ':' not in x, re.findall('(?<=\[\[)[^\]]*(?=\]\])',elem.text))
                for l in links:
                    l = re.split('[#|]',l.lower())[0]
                    if l in self.titles:
                        self.transitions[docid][self.titles[l]] = 1
                        ones += 1
            elem.clear()
            
        for i in xrange(len(self.transitions)):
            if 1 not in self.transitions[i]:
                self.transitions[i].fill(1/N)
            else:
                os = self.transitions[i].tolist().count(1)
                self.transitions[i] = self.transitions[i]*(1-ALPHA)/os + ALPHA/N
            smr = sum(self.transitions[i])
            assert smr < 1.001 and smr > 0.999

    def pagerank(self):
        vector = random.choice(self.transitions)
        for i in xrange(ITERATIONS):
            vector = dot(vector, self.transitions)
        f = open(self._file[:-4]+'_pagerank.txt','w')
        for i in vector:
            f.write(str(i)+'\n')
        f.close()
        

if __name__ == "__main__":
    arg = sys.argv[1:]
    if len(arg) != 1:
        print '''Usage: pagerank <collection.xml>'''
    file = arg[0]
    try:
        import psyco
        psyco.full()
    except:
        pass
    pr = PageRank(file)
    
    